#include "kernel_operator.h"
#include <type_traits>
using namespace AscendC;
constexpr int32_t BUFFER_NUM = 2;

template<typename TYPE_Y> class KernelRmsNorm {
    using T = TYPE_Y;
public:
    __aicore__ inline KernelRmsNorm() {}
    __aicore__ inline void Init(GM_ADDR x1, GM_ADDR x2, GM_ADDR y, GM_ADDR rstd,
                            uint32_t loop_num, uint32_t x2_size, float epsilon, uint32_t finalTileNum, uint32_t tileDataNum, uint32_t TailDataNum) {
        ASSERT(GetBlockNum() != 0 && "block dim can not be zero!");
        this->loop_num = loop_num;
        this->x2_size = x2_size;
        this->epsilon = epsilon;

        this->tileNum = finalTileNum;
        this->tileDataNum = tileDataNum;
        this->tailDataNum = TailDataNum;

        x1Gm.SetGlobalBuffer((__gm__ DTYPE_X1*)x1, 1);
        x2Gm.SetGlobalBuffer((__gm__ DTYPE_X2*)x2, this->x2_size);
        yGm.SetGlobalBuffer((__gm__ DTYPE_Y*)y, 1);
        rstdGm.SetGlobalBuffer((__gm__ DTYPE_RSTD*)rstd, this->loop_num);

        pipe.InitBuffer(inQueueX1, BUFFER_NUM, this->tileDataNum * sizeof(DTYPE_X1));
        pipe.InitBuffer(inQueueX2, BUFFER_NUM, this->tileDataNum * sizeof(DTYPE_X2));
        pipe.InitBuffer(outQueueY, BUFFER_NUM, this->tileDataNum * sizeof(DTYPE_Y));
        
        if constexpr (std::is_same_v<T, float>) {
        }
        else
        {
            pipe.InitBuffer(QueueTmp1, this->tileDataNum * sizeof(float));
            pipe.InitBuffer(QueueTmp2, this->tileDataNum * sizeof(float));
        }
        pipe.InitBuffer(QueueTmp3, 64* sizeof(float));
    }
    __aicore__ inline void Process() {
        auto tmp_sum = QueueTmp3.Get<float>();
        tmp_sum.SetValue(16, (float)(this->x2_size));
        tmp_sum.SetValue(32, (float)(1));

        if constexpr (std::is_same_v<T, float>) {
            Sqrt(tmp_sum[16], tmp_sum[16], 1);
        }

        float rstd_l;

        for(int32_t i=0; i<this->loop_num; i++)
        {
            Duplicate<float>(tmp_sum, 0, 1);
            int32_t loopCount = this->tileNum;
            this->processDataNum = this->tileDataNum;
            for (int32_t j = 0; j < loopCount; j++) {
                if (j == this->tileNum - 1) {
                    this->processDataNum = this->tailDataNum;
                }
                CopyIn1(i, j);
                Compute1(tmp_sum);
            }
            
            if constexpr (std::is_same_v<T, float>) {

                Sqrt(tmp_sum, tmp_sum, 1);
                Div(tmp_sum, tmp_sum[16], tmp_sum, 1);

                rstd_l = tmp_sum.GetValue(0);

                rstdGm.SetValue(i, rstd_l);
            }
            else {

                Div(tmp_sum, tmp_sum, tmp_sum[16], 1);
                

                Cast(tmp_sum.ReinterpretCast<DTYPE_X1>(), tmp_sum, RoundMode::CAST_RINT, 1);
                Cast(tmp_sum, tmp_sum.ReinterpretCast<DTYPE_X1>(), RoundMode::CAST_NONE, 1);

                Sqrt(tmp_sum, tmp_sum, 1);
                Cast(tmp_sum.ReinterpretCast<DTYPE_X1>(), tmp_sum, RoundMode::CAST_RINT, 1);
                Cast(tmp_sum, tmp_sum.ReinterpretCast<DTYPE_X1>(), RoundMode::CAST_NONE, 1);
                Div(tmp_sum, tmp_sum[32], tmp_sum, 1);
                
                Cast(tmp_sum.ReinterpretCast<DTYPE_RSTD>(), tmp_sum, RoundMode::CAST_RINT, 1);
                DTYPE_RSTD rstd_s = tmp_sum.ReinterpretCast<DTYPE_RSTD>().GetValue(0);

                if constexpr (std::is_same_v<T, half>) {
                    rstd_l = (float)rstd_s;
                }
                else {
                    rstd_l = ToFloat(rstd_s);
                }

                rstdGm.SetValue(i, rstd_s);
            }

            this->processDataNum = this->tileDataNum;
            for (int32_t j = 0; j < loopCount; j++) {
                if (j == this->tileNum - 1) {
                    this->processDataNum = this->tailDataNum;
                }
                CopyIn2(i, j);
                Compute2(rstd_l);
                CopyOut2(i, j);
            }
            
        }
    }
    private:
    __aicore__ inline void CopyIn1(int32_t progress, int32_t num) {
        LocalTensor<DTYPE_X1> x1Local = inQueueX1.AllocTensor<DTYPE_X1>();
        DataCopy(x1Local, x1Gm[progress * this->x2_size + num * this->tileDataNum], this->processDataNum);
        inQueueX1.EnQue(x1Local);
    }
    __aicore__ inline void Compute1(LocalTensor<float> sum) {
        LocalTensor<DTYPE_X1> x1Local = inQueueX1.DeQue<DTYPE_X1>();
        

        if constexpr (std::is_same_v<T, float>) {
            // auto tmp1 = QueueTmp1.Get<float>();
            // auto tmp2 = QueueTmp2.Get<T>();
            LocalTensor<DTYPE_X2> x2Local = inQueueX2.AllocTensor<DTYPE_X2>();

            Mul(x1Local, x1Local, x1Local, this->processDataNum);
            ReduceSum(x1Local, x1Local, x2Local, this->processDataNum);

            Add(sum, sum, x1Local.template ReinterpretCast<float>(), 1);

            inQueueX2.FreeTensor(x2Local);
        }
        else
        {
            auto tmp1 = QueueTmp1.Get<float>();
            auto tmp2 = QueueTmp2.Get<T>();

            Cast(tmp1, x1Local, RoundMode::CAST_NONE, this->processDataNum);
            Mul(tmp1, tmp1, tmp1, this->processDataNum);

            Cast(tmp2, tmp1, RoundMode::CAST_NONE, this->processDataNum);
            Cast(tmp1, tmp2, RoundMode::CAST_RINT, this->processDataNum);

            ReduceSum(tmp1, tmp1, tmp2.template ReinterpretCast<float>(), this->processDataNum);
            Add(sum, sum, tmp1, 1);
        }

        inQueueX1.FreeTensor(x1Local);
    }
    __aicore__ inline void CopyIn2(int32_t progress, int32_t num) {
        LocalTensor<DTYPE_X1> x1Local = inQueueX1.AllocTensor<DTYPE_X1>();
        LocalTensor<DTYPE_X2> x2Local = inQueueX2.AllocTensor<DTYPE_X2>();
        DataCopy(x1Local, x1Gm[progress * this->x2_size + num * this->tileDataNum], this->processDataNum);
        DataCopy(x2Local, x2Gm[num * this->tileDataNum], this->processDataNum);
        inQueueX1.EnQue(x1Local);
        inQueueX2.EnQue(x2Local);
    }
    
    __aicore__ inline void Compute2(float rstd) {
        LocalTensor<DTYPE_X1> x1Local = inQueueX1.DeQue<DTYPE_X1>();
        LocalTensor<DTYPE_X2> x2Local = inQueueX2.DeQue<DTYPE_X2>();
        LocalTensor<DTYPE_Y> yLocal = outQueueY.AllocTensor<DTYPE_Y>();
        
        if constexpr (std::is_same_v<T, float>) {
            Mul(x1Local, x1Local, x2Local, this->processDataNum);
            Muls(yLocal, x1Local, rstd, this->processDataNum);
        }
        else
        {
            auto tmp1 = QueueTmp1.Get<float>();
            auto tmp2 = QueueTmp2.Get<float>();

            Cast(tmp1, x1Local, RoundMode::CAST_NONE, this->processDataNum);
            Muls(tmp1, tmp1, rstd, this->processDataNum);

            Cast(tmp2.ReinterpretCast<DTYPE_X1>(), tmp1, RoundMode::CAST_RINT, this->processDataNum);
            Cast(tmp1, tmp2.ReinterpretCast<DTYPE_X1>(), RoundMode::CAST_NONE, this->processDataNum);

            Cast(tmp2, x2Local, RoundMode::CAST_NONE, this->processDataNum);
            Mul(tmp1, tmp1, tmp2, this->processDataNum);
            
            Cast(yLocal, tmp1, RoundMode::CAST_RINT, this->processDataNum);
        }
        inQueueX1.FreeTensor(x1Local);
        inQueueX2.FreeTensor(x2Local);
        outQueueY.EnQue<TYPE_Y>(yLocal);
    }

    __aicore__ inline void CopyOut2(int32_t progress, int32_t num)  {
        LocalTensor<DTYPE_Y> yLocal = outQueueY.DeQue<DTYPE_Y>();
        DataCopy(yGm[progress * this->x2_size + num * this->tileDataNum], yLocal, this->processDataNum);
        outQueueY.FreeTensor(yLocal);
    }
private:
    TPipe pipe;
    TQue<QuePosition::VECIN, BUFFER_NUM> inQueueX1, inQueueX2;
    TQue<QuePosition::VECOUT, BUFFER_NUM> outQueueY;
    TBuf<QuePosition::VECCALC> QueueTmp1, QueueTmp2, QueueTmp3, QueueTmp4, boox;

    GlobalTensor<DTYPE_X1> x1Gm;
    GlobalTensor<DTYPE_X2> x2Gm;
    GlobalTensor<DTYPE_Y> yGm;
    GlobalTensor<DTYPE_RSTD> rstdGm;
    int32_t loop_num;
    int32_t x2_size;
    float epsilon;

    int32_t tileNum;
    int32_t tileDataNum;
    int32_t tailDataNum;
    int32_t processDataNum;
};

template<typename TYPE_Y> class KernelRmsNorm_all {
    using T = TYPE_Y;
public:
    __aicore__ inline KernelRmsNorm_all() {}
    __aicore__ inline void Init(GM_ADDR x1, GM_ADDR x2, GM_ADDR y, GM_ADDR rstd,
                            uint32_t loop_num, uint32_t x2_size, float epsilon, uint32_t finalTileNum, uint32_t tileDataNum, uint32_t TailDataNum) {
        ASSERT(GetBlockNum() != 0 && "block dim can not be zero!");
        this->loop_num = loop_num;
        this->x2_size = x2_size;
        this->epsilon = epsilon;

        this->tileNum = finalTileNum;
        this->tileDataNum = tileDataNum;
        this->tailDataNum = TailDataNum;

        x1Gm.SetGlobalBuffer((__gm__ DTYPE_X1*)x1, 1);
        x2Gm.SetGlobalBuffer((__gm__ DTYPE_X2*)x2, this->x2_size);
        yGm.SetGlobalBuffer((__gm__ DTYPE_Y*)y, 1);
        rstdGm.SetGlobalBuffer((__gm__ DTYPE_RSTD*)rstd, this->loop_num);

        pipe.InitBuffer(inQueueX1, BUFFER_NUM, this->tileDataNum * sizeof(DTYPE_X1));
        pipe.InitBuffer(inQueueX2, BUFFER_NUM, this->tileDataNum * sizeof(DTYPE_X2));
        pipe.InitBuffer(outQueueY, BUFFER_NUM, this->tileDataNum * sizeof(DTYPE_Y));

        pipe.InitBuffer(QueueTmp1, this->tileDataNum * sizeof(float));
        pipe.InitBuffer(QueueTmp2, this->tileDataNum * sizeof(float));
        pipe.InitBuffer(QueueTmp3, 64* sizeof(float));
    }
    __aicore__ inline void Process() {
        auto tmp_sum = QueueTmp3.Get<float>();
        tmp_sum.SetValue(16, (float)(this->x2_size));
        tmp_sum.SetValue(32, (float)(1));
        if constexpr (std::is_same_v<T, float>) {
            Sqrt(tmp_sum[16], tmp_sum[16], 1);
        }

        float rstd_l;

        for(int32_t i=0; i<this->loop_num; i++)
        {
            Duplicate<float>(tmp_sum, 0, 1);
            int32_t loopCount = this->tileNum;
            this->processDataNum = this->tileDataNum;
            for (int32_t j = 0; j < loopCount; j++) {
                if (j == this->tileNum - 1) {
                this->processDataNum = this->tailDataNum;
                }
                CopyIn1(i, j);
                Compute1(tmp_sum);
            }
            if constexpr (std::is_same_v<T, float>) {

                Sqrt(tmp_sum, tmp_sum, 1);
                Div(tmp_sum, tmp_sum[16], tmp_sum, 1);

                rstd_l = tmp_sum.GetValue(0);

                rstdGm.SetValue(i, rstd_l);
            }
            else {

                Div(tmp_sum, tmp_sum, tmp_sum[16], 1);
                

                Cast(tmp_sum.ReinterpretCast<DTYPE_X1>(), tmp_sum, RoundMode::CAST_RINT, 1);
                Cast(tmp_sum, tmp_sum.ReinterpretCast<DTYPE_X1>(), RoundMode::CAST_NONE, 1);
                Sqrt(tmp_sum, tmp_sum, 1);
                Cast(tmp_sum.ReinterpretCast<DTYPE_X1>(), tmp_sum, RoundMode::CAST_RINT, 1);
                Cast(tmp_sum, tmp_sum.ReinterpretCast<DTYPE_X1>(), RoundMode::CAST_NONE, 1);
                Div(tmp_sum, tmp_sum[32], tmp_sum, 1);
                Cast(tmp_sum.ReinterpretCast<DTYPE_RSTD>(), tmp_sum, RoundMode::CAST_RINT, 1);
                DTYPE_RSTD rstd_s = tmp_sum.ReinterpretCast<DTYPE_RSTD>().GetValue(0);
                if constexpr (std::is_same_v<T, half>) {
                    rstd_l = (float)rstd_s;
                }
                else {
                    rstd_l = ToFloat(rstd_s);
                }
                rstdGm.SetValue(i, rstd_s);
            }

            this->processDataNum = this->tileDataNum;
            for (int32_t j = 0; j < loopCount; j++) {
                if (j == this->tileNum - 1) {
                this->processDataNum = this->tailDataNum;
                }
                CopyIn2(i, j);
                Compute2(rstd_l);
                CopyOut2(i, j);
            }
            
        }
    }
    private:
    __aicore__ inline void CopyIn1(int32_t progress, int32_t num) {
        LocalTensor<DTYPE_X1> x1Local = inQueueX1.AllocTensor<DTYPE_X1>();
        DataCopyExtParams copyParams{1, static_cast<uint32_t>(this->processDataNum * sizeof(DTYPE_X1)), 0, 0, 0};
        DataCopyPadExtParams<DTYPE_X1> padParams{false, 0, 0, 0};
        DataCopyPad(x1Local, x1Gm[progress * this->x2_size + num * this->tileDataNum], copyParams, padParams);
        inQueueX1.EnQue(x1Local);
    }
    __aicore__ inline void Compute1(LocalTensor<float> sum) {
        LocalTensor<DTYPE_X1> x1Local = inQueueX1.DeQue<DTYPE_X1>();
        auto tmp1 = QueueTmp1.Get<float>();
        auto tmp2 = QueueTmp2.Get<T>();

        if constexpr (std::is_same_v<T, float>) {
            Mul(tmp1, x1Local, x1Local, this->processDataNum);
            ReduceSum(tmp1, tmp1, tmp2.template ReinterpretCast<float>(), this->processDataNum);

            Add(sum, sum, tmp1, 1);
        }
        else
        {
            Cast(tmp1, x1Local, RoundMode::CAST_NONE, this->processDataNum);
            Mul(tmp1, tmp1, tmp1, this->processDataNum);

            Cast(tmp2, tmp1, RoundMode::CAST_NONE, this->processDataNum);
            Cast(tmp1, tmp2, RoundMode::CAST_RINT, this->processDataNum);

            ReduceSum(tmp1, tmp1, tmp2.template ReinterpretCast<float>(), this->processDataNum);
            Add(sum, sum, tmp1, 1);
        }

        inQueueX1.FreeTensor(x1Local);
    }
    __aicore__ inline void CopyIn2(int32_t progress, int32_t num) {

        LocalTensor<DTYPE_X1> x1Local = inQueueX1.AllocTensor<DTYPE_X1>();
        LocalTensor<DTYPE_X2> x2Local = inQueueX2.AllocTensor<DTYPE_X2>();
        DataCopyExtParams copyParams{1, static_cast<uint32_t>(this->processDataNum * sizeof(DTYPE_X1)), 0, 0, 0};
        DataCopyPadExtParams<DTYPE_X1> padParams{false, 0, 0, 0};
        DataCopyPad(x1Local, x1Gm[progress * this->x2_size + num * this->tileDataNum], copyParams, padParams);
        DataCopyPad(x2Local, x2Gm[num * this->tileDataNum], copyParams, padParams);
        inQueueX1.EnQue(x1Local);
        inQueueX2.EnQue(x2Local);
    }
    
    __aicore__ inline void Compute2(float rstd) {
        LocalTensor<DTYPE_X1> x1Local = inQueueX1.DeQue<DTYPE_X1>();
        LocalTensor<DTYPE_X2> x2Local = inQueueX2.DeQue<DTYPE_X2>();
        LocalTensor<DTYPE_Y> yLocal = outQueueY.AllocTensor<DTYPE_Y>();
        
        if constexpr (std::is_same_v<T, float>) {
            Mul(x1Local, x1Local, x2Local, this->processDataNum);
            Muls(yLocal, x1Local, rstd, this->processDataNum);
        }
        else
        {
            auto tmp1 = QueueTmp1.Get<float>();
            auto tmp2 = QueueTmp2.Get<float>();

            Cast(tmp1, x1Local, RoundMode::CAST_NONE, this->processDataNum);
            Muls(tmp1, tmp1, rstd, this->processDataNum);

            Cast(tmp2.ReinterpretCast<DTYPE_X1>(), tmp1, RoundMode::CAST_RINT, this->processDataNum);
            Cast(tmp1, tmp2.ReinterpretCast<DTYPE_X1>(), RoundMode::CAST_NONE, this->processDataNum);

            Cast(tmp2, x2Local, RoundMode::CAST_NONE, this->processDataNum);
            Mul(tmp1, tmp1, tmp2, this->processDataNum);
            
            Cast(yLocal, tmp1, RoundMode::CAST_RINT, this->processDataNum);
        }
        inQueueX1.FreeTensor(x1Local);
        inQueueX2.FreeTensor(x2Local);
        outQueueY.EnQue<TYPE_Y>(yLocal);
    }

    __aicore__ inline void CopyOut2(int32_t progress, int32_t num)  {
        LocalTensor<DTYPE_Y> yLocal = outQueueY.DeQue<DTYPE_Y>();
        DataCopyExtParams copyParams{1, static_cast<uint32_t>(this->processDataNum * sizeof(DTYPE_Y)), 0, 0, 0};
        DataCopyPad(yGm[progress * this->x2_size + num * this->tileDataNum], yLocal, copyParams);
        outQueueY.FreeTensor(yLocal);
    }
private:
    TPipe pipe;
    TQue<QuePosition::VECIN, BUFFER_NUM> inQueueX1, inQueueX2;
    TQue<QuePosition::VECOUT, BUFFER_NUM> outQueueY;
    TBuf<QuePosition::VECCALC> QueueTmp1, QueueTmp2, QueueTmp3, QueueTmp4, boox;

    GlobalTensor<DTYPE_X1> x1Gm;
    GlobalTensor<DTYPE_X2> x2Gm;
    GlobalTensor<DTYPE_Y> yGm;
    GlobalTensor<DTYPE_RSTD> rstdGm;
    int32_t loop_num;
    int32_t x2_size;
    float epsilon;

    int32_t tileNum;
    int32_t tileDataNum;
    int32_t tailDataNum;
    int32_t processDataNum;
};
extern "C" __global__ __aicore__ void rms_norm(GM_ADDR x1, GM_ADDR x2, GM_ADDR y, GM_ADDR rstd, GM_ADDR workspace, GM_ADDR tiling) {
    GET_TILING_DATA(tiling_data, tiling);
    // TODO: user kernel impl

    if (TILING_KEY_IS(1)) {
        KernelRmsNorm<DTYPE_X1> op;
        op.Init(x1, x2, y, rstd,
            tiling_data.loop_num, tiling_data.x2_size, tiling_data.epsilon, tiling_data.finalTileNum, tiling_data.tileDataNum, tiling_data.TailDataNum);
        op.Process();
    }
    else if (TILING_KEY_IS(2)) {
        KernelRmsNorm_all<DTYPE_X1> op;
        op.Init(x1, x2, y, rstd,
            tiling_data.loop_num, tiling_data.x2_size, tiling_data.epsilon, tiling_data.finalTileNum, tiling_data.tileDataNum, tiling_data.TailDataNum);
        op.Process();
    }
    
}